{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Misc imports\n",
    "import os\n",
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "# Import Minions + Minions\n",
    "from minions.minions import Minions\n",
    "from minions.minion import Minion\n",
    "\n",
    "# Import Minion Clients\n",
    "from minions.clients.ollama import OllamaClient\n",
    "from minions.clients.tokasaurus import TokasaurusClient\n",
    "from minions.clients.openai import OpenAIClient\n",
    "from minions.clients.anthropic import AnthropicClient\n",
    "from minions.clients.together import TogetherClient\n",
    "\n",
    "# Import Pydantic\n",
    "from pydantic import BaseModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure Clients\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Specify structured output schema for the OllamaClient. Run this block as is! Do _NOT_ make any modifications\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StructuredLocalOutput(BaseModel):\n",
    "    explanation: str\n",
    "    citation: str | None\n",
    "    answer: str | None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Instantiate Clients: Here we instantiate our local client to be ollama and remote client to be OpenAI\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting server with command: tksrs             port=57475             model=meta-llama/Llama-3.2-1B-Instruct             torch_compile=T             \n",
      "Waiting for server to start on port 57475...\n",
      "ConnectionError: HTTPConnectionPool(host='localhost', port=57475): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7fd732122e10>: Failed to establish a new connection: [Errno 111] Connection refused'))\n",
      "Server not yet started (attempt 0) retrying... (this is normal)\n",
      "ConnectionError: HTTPConnectionPool(host='localhost', port=57475): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7fd738635640>: Failed to establish a new connection: [Errno 111] Connection refused'))\n",
      "Server not yet started (attempt 1) retrying... (this is normal)\n",
      "Setting max_num_tokens_per_request to 131072\n",
      "Starting 3 processes: ['model_worker', 'manager', 'server']\n",
      "Running in the main process: server\n",
      "ConnectionError: HTTPConnectionPool(host='localhost', port=57475): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7fd73216a7b0>: Failed to establish a new connection: [Errno 111] Connection refused'))\n",
      "Server not yet started (attempt 2) retrying... (this is normal)\n",
      "2025-02-25 06:39:43 | INFO | server | Starting web server\n",
      "ConnectionError: HTTPConnectionPool(host='localhost', port=57475): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7fd73216bda0>: Failed to establish a new connection: [Errno 111] Connection refused'))\n",
      "Server not yet started (attempt 3) retrying... (this is normal)\n",
      "2025-02-25 06:39:46 | INFO | model_worker | Model worker started!\n",
      "2025-02-25 06:39:46 | INFO | model_worker | Creating model on device cuda:0 with dtype torch.bfloat16\n",
      "ConnectionError: HTTPConnectionPool(host='localhost', port=57475): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7fd732194620>: Failed to establish a new connection: [Errno 111] Connection refused'))\n",
      "Server not yet started (attempt 4) retrying... (this is normal)\n",
      "Building layers 0 to 16\n",
      "Loading from safetensors\n",
      "2025-02-25 06:39:47 | INFO | manager | Manager started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fetching 7 files: 100%|██████████| 7/7 [00:00<00:00, 104113.93it/s]\n",
      "Loading safetensors files: 100%|██████████| 1/1 [00:00<00:00,  2.82it/s]\n",
      "INFO:     Started server process [27630]\n",
      "INFO:     Waiting for application startup.\n",
      "INFO:     Application startup complete.\n",
      "INFO:     Uvicorn running on http://0.0.0.0:57475 (Press CTRL+C to quit)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-02-25 06:39:47 | INFO | model_worker | Created model\n",
      "System startup time: 4.850164175033569\n",
      "2025-02-25 06:39:47 | INFO | model_worker | Model loop started!\n",
      "              _              _                       \n",
      "  ___   __ _ | |_       ___ | |  ___   ___  _ __     \n",
      " / _ \\ / _` || __|     / __|| | / _ \\ / _ \\| '_ \\    \n",
      "|  __/| (_| || |_  _   \\__ \\| ||  __/|  __/| |_) | _ \n",
      " \\___| \\__,_| \\__|( )  |___/|_| \\___| \\___|| .__/ ( )\n",
      "                  |/                       |_|    |/ \n",
      " _          __                                                                           _      \n",
      "(_) _ __   / _|  ___  _ __   ___  _ __    ___   ___      _ __   ___  _ __    ___   __ _ | |_    \n",
      "| || '_ \\ | |_  / _ \\| '__| / _ \\| '_ \\  / __| / _ \\    | '__| / _ \\| '_ \\  / _ \\ / _` || __|   \n",
      "| || | | ||  _||  __/| |   |  __/| | | || (__ |  __/ _  | |   |  __/| |_) ||  __/| (_| || |_  _ \n",
      "|_||_| |_||_|   \\___||_|    \\___||_| |_| \\___| \\___|( ) |_|    \\___|| .__/  \\___| \\__,_| \\__|(_)\n",
      "                                                    |/              |_|                         \n",
      "\n",
      "2025-02-25 06:39:47 | INFO | manager | Submitting max-sized batch...\n",
      "INFO:     127.0.0.1:42950 - \"GET /health HTTP/1.1\" 404 Not Found\n",
      "Started server with pid 27629\n"
     ]
    }
   ],
   "source": [
    "remote_client = OpenAIClient(model_name=\"gpt-4o\", temperature=0.0)\n",
    "\n",
    "# Option 1: Ollama\n",
    "local_client = OllamaClient(\n",
    "    model_name=\"llama3.2:3b\", \n",
    "    temperature=0.0, \n",
    "    structured_output_schema=StructuredLocalOutput\n",
    ")\n",
    "\n",
    "# option 2\n",
    "# local_client = TokasaurusClient(\n",
    "#     model_name=\"meta-llama/Llama-3.2-1B-Instruct\",\n",
    "#     temperature=0.0,\n",
    "#     # Structured outputs are not yet supported with tokasaurus\n",
    "#     # structured_output_schema=StructuredLocalOutput\n",
    "# )\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set-up Communication Protocol\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option 1: Minions\n",
    "# protocol = Minions(local_client=local_client, remote_client=remote_client)\n",
    "\n",
    "# Option 2: Minion\n",
    "protocol = Minion(local_client=local_client, remote_client=remote_client)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure Minion/Minions Task\n",
    "\n",
    "- context (List[str]): context that the minions need to reason over\n",
    "- doc_metadata (str): every task is parameterized by doc_metadata which describes the \"type\" of information contained in the context\n",
    "- task (str): description of the query to be completed over the context\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Specify Input Context\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"\"\"\\\n",
    "Patient Name: John A. Doe\n",
    "Medical Record Number: 001234567\n",
    "DOB: 1967-08-22\n",
    "Gender: Male\n",
    "Admission Date: 2025-01-15\n",
    "Discharge Date: 2025-01-20\n",
    "\n",
    "Chief Complaint:\n",
    "The patient presented with chest discomfort, shortness of breath, and fatigue. Symptoms began gradually 48 hours before admission and intensified over time.\n",
    "\n",
    "History of Present Illness:\n",
    "John A. Doe, a 57-year-old male with a history of type 2 diabetes mellitus and hypertension, was admitted following several episodes of atypical chest pain. He experienced intermittent chest tightness and dyspnea on exertion. The pain did not radiate, and there was no associated nausea.\n",
    "\n",
    "Past Medical History:\n",
    "- Type 2 Diabetes Mellitus (diagnosed 2010) – managed with metformin.\n",
    "- Hypertension (diagnosed 2012) – treated with lisinopril.\n",
    "- Hyperlipidemia – treated with atorvastatin.\n",
    "- Former smoker (quit 2015).\n",
    "- No known drug allergies.\n",
    "\n",
    "Medications on Admission:\n",
    "- Metformin 1000 mg twice daily.\n",
    "- Lisinopril 20 mg daily.\n",
    "- Atorvastatin 40 mg nightly.\n",
    "- Aspirin 81 mg daily.\n",
    "\n",
    "Physical Examination:\n",
    "General: The patient is alert and oriented, in moderate distress.\n",
    "Cardiovascular: Blood pressure 150/95 mmHg; heart rate 88 bpm; regular rhythm.\n",
    "Respiratory: Lungs are clear to auscultation bilaterally; respiratory rate 16 bpm.\n",
    "Neurological: No focal deficits observed.\n",
    "\n",
    "Laboratory and Diagnostic Findings (Admission):\n",
    "- Complete Blood Count: WBC 7.2 x10^9/L, Hemoglobin 13.5 g/dL.\n",
    "- Basic Metabolic Panel: Serum creatinine 1.0 mg/dL; electrolytes normal.\n",
    "- Fasting Blood Glucose: 180 mg/dL.\n",
    "- Hemoglobin A1C: 8.5%.\n",
    "- Lipid Profile: Total Cholesterol 220 mg/dL, LDL 140 mg/dL, HDL 35 mg/dL.\n",
    "- ECG: Non-specific ST-T changes.\n",
    "- Chest X-ray: No acute findings.\n",
    "\n",
    "Progress Notes:\n",
    "Day 2: Initiated subcutaneous insulin to complement oral hypoglycemics. Blood pressure remained elevated at 148/92 mmHg.\n",
    "Day 3: Follow-up labs showed a minor improvement in glycemic levels and stable kidney function.\n",
    "Day 5: Vital signs improved with a blood pressure of 140/85 mmHg and heart rate of 80 bpm. A repeat lab indicated that the fasting blood glucose had decreased significantly compared to admission.\n",
    "Day 6: Endocrinology recommended maintaining current treatment with a possible gradual reduction in insulin if improvements continue.\n",
    "Day 7: Patient stabilized and was discharged with clear instructions for follow-up with both primary care and endocrinology.\n",
    "\n",
    "Discharge Summary:\n",
    "John A. Doe was discharged on Day 7 after noticeable improvements in glycemic control and blood pressure. The patient is to continue current medications with a scheduled outpatient review.\n",
    "Medications at Discharge:\n",
    "- Metformin 1000 mg BID.\n",
    "- Lisinopril 20 mg daily.\n",
    "- Atorvastatin 40 mg nightly.\n",
    "- Adjusted insulin regimen based on recent glycemic trends.\n",
    "\n",
    "Additional Information:\n",
    "- Admission weight: 95 kg (BMI: 31.2).\n",
    "- Dietary consult completed; patient advised on a diabetic diet.\n",
    "- Patient educated on adherence and regular blood glucose monitoring.\n",
    "\n",
    "Summary:\n",
    "This electronic health record details John A. Doe's hospital admission, clinical findings, laboratory values, treatments, and progress over a 7-day period. Key values include a fasting blood glucose of 180 mg/dL on admission and 150 mg/dL on Day 5, indicating a notable improvement in glycemic control.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Specify Metadata and Task Description\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_metadata = \"Patient Visit Notes\"\n",
    "\n",
    "# Define the task without explicitly providing the key values.\n",
    "task = (\n",
    "    \"Using the electronic health records provided, calculate the percentage decrease in the patient's fasting \"\n",
    "    \"blood glucose level from admission to Day 5. Extract the necessary values from the record, show your calculations, \"\n",
    "    \"and provide the final percentage decrease.\"\n",
    ")\n",
    "\n",
    "# For testing purposes, we suggest that the correct answer is:\n",
    "suggested_final_answer = (\n",
    "    \"16.67% decrease in fasting blood glucose from admission to Day 5.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the protocol\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "entering loop with max_rounds:  5\n",
      "getting worker's response\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/lib/python3.12/site-packages/torch/_inductor/compile_fx.py:194: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-02-25 06:40:01 | INFO | manager | ... Max-sized batch complete in 13.61s\n",
      "2025-02-25 06:40:01 | INFO | manager | Submitting more batches to try to trigger recompiles...\n",
      "2025-02-25 06:40:37 | INFO | manager | ... Compilation batches complete in 36.55s\n",
      "2025-02-25 06:40:37 | INFO | manager | Starting manager loop\n",
      "INFO:     127.0.0.1:50792 - \"POST /v1/chat/completions HTTP/1.1\" 200 OK\n"
     ]
    }
   ],
   "source": [
    "output = protocol(\n",
    "        task=task,\n",
    "        doc_metadata=doc_metadata,\n",
    "        context=[context],\n",
    "        max_rounds=5,  # you can adjust rounds as needed for testing\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The patient's fasting blood glucose level decreased by 16.7% from admission to Day 5.\n"
     ]
    }
   ],
   "source": [
    "print(output['final_answer'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "minions",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
